#pragma once
#include <mutex>
#include <opencv2/core.hpp>
#include <opencv2/highgui.hpp>
#include <string>
#include <thread>
#ifdef USE_ROS
// ROS
#include <cv_bridge/cv_bridge.h>
#include <image_transport/image_transport.h>
#include <ros/ros.h>
#endif

#include "udp.hpp"
// FFmpeg
extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/avutil.h>
#include <libavutil/pixdesc.h>
#include <libswscale/swscale.h>
}
const int NORMAL_STREAM = 0;
const int UDP_STREAM = 1;
const int MANUAL_IO = 2;

class LibAV {
 public:
  LibAV(int flag = NORMAL_STREAM, std::string topic = "image");
#ifdef USE_ROS
  LibAV(ros::NodeHandle* nh, int flag = NORMAL_STREAM,
        std::string topic = "image");
#endif
  ~LibAV();
  /**
   * @brief Initialize for decoder.
   *
   * @param src - Source address of video stream
   * @param min_frame_buffer - Only used in io mode. Minimum buffer for decoding
   * single frame. Need to be set properly.
   * Too large will cause large frame delay.
   * Too small may enough for 1 frame.
   * @param buffer_size - Buffer size for normal av decode. Or size of template
   * buffer for io mode.
   *
   * @return Whether successfully initialize
   */
  bool init(std::string src = "", int min_frame_buffer = 100000,
            int buffer_size = 262144);
  cv::Mat getMat();
  inline bool endOfStream() { return end_of_stream_; }
  void saveToBuffer(std::vector<uint8_t>& buf, int offset = 0);
  void saveToBuffer(uint8_t* buf, int length, int offset = 0);

 private:
  inline void printBufferSize() {
    std::cout << "\033[36mbuffer size:" << frame_buffer_.size() << " \033[0m"
              << std::endl;
  }
  inline void waitForBuffer(int min_size, int step_time = 50) {
    while (frame_buffer_.size() < min_size) {
      std::cout << "waiting for data: " << frame_buffer_.size() << "/"
                << min_size << std::endl;
      std::this_thread::sleep_for(std::chrono::milliseconds(step_time));
    }
  }
  int mode_;
  std::mutex image_lock_;
  std::thread* stream_th_;

  AVFormatContext* inctx_ = nullptr;
  int vstrm_idx_;
  AVStream* vstrm_;
  AVFrame* frame_;
  AVFrame* decframe_;
  SwsContext* swsctx_;

  int dst_width_, dst_height_;
  std::vector<uint8_t> framebuf_;
  unsigned int nb_frames_ = 0;
  bool end_of_stream_ = false;
  int got_pic_ = 0;
  AVPacket pkt_;
  void openThread();

#ifdef USE_ROS
  // ROS
  ros::NodeHandle* nh_;
  image_transport::ImageTransport* it_;
  image_transport::Publisher img_pub_;
#endif
  int initAV(std::string& src, int buffer_size);

  static int readPacket(void* opaque, uint8_t* buf, int buf_size);
  AVIOContext* avio_ctx_;
  int initIOAV(int min_frame_buffer, int buffer_size);
  uint8_t* io_buffer_;

  UDP udp_;
  std::thread* udp_t_;
  std::mutex buf_lock_;
  std::vector<uint8_t> frame_buffer_;
  int min_frame_buffer_;
  int groupUDP(std::string& uri);
};
